import os
import time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision as tv
import torchvision.transforms as transforms
import torchnet as tnt
from args import *
from tool import *
from wheel import *
import evaluate
from models import *


save_path = os.path.join(args.log_path, "ablation")
if not os.path.exists(save_path):
    os.makedirs(save_path)
log_file_path = os.path.join(save_path, "log.{}".format(timestamp()))
log_file = open(log_file_path, "a")
log_file.write("begin time: {}\n".format(time.asctime()))


test_loader = torch.utils.data.DataLoader(
    tv.datasets.MNIST(args.data_path, train=False,
                   transform=transforms.Compose([
                    #    transforms.Resize([224, 224]),
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,)),
                    #    transforms.Lambda(lambda x: x.repeat(3, 1, 1))
                   ])),
    batch_size=args.batch_size, shuffle=True)


def test(model, model_type):
    model.eval()
    n_correct = 0
    with torch.no_grad():
        for image, label in test_loader:
            X, L = image.cuda(), label.cuda()
            fea, logit = model(X)

            if model_type == "a-softmax":
                logit = logit[0]  # logit = (cos(x), phi(x))
            elif model_type in ("c-contrastive", "c-triplet"):
                logit = - logit  # euc dist -> smaller is more similar
            pred = logit.argmax(dim=1).cpu()
            n_correct += pred.eq(label.view(-1)).sum().item()

    acc = n_correct / len(test_loader.dataset)

    return acc

T_list = ["naive", "normface", "a-softmax", "am-softmax", "centre"]

P_list = [
    "naive_0bias_0.0lambda_20scale_1dy-scale",
    "normface_0bias_0.0lambda_20.0scale_1dy-scale",
    "a-softmax_0bias_0.0lambda_20scale_1dy-scale",
    "am-softmax_0.4m_5.0s",
    "centre_0bias_1.0lambda_20scale_1dy-scale"
]
P_list = [os.path.join(args.log_path, _p) for _p in P_list]

M_list = [
    "26e_0.9869acc.pth",
    "25e_0.994acc.pth",
    "23e_0.9917acc.pth",
    "19e_0.9916acc.pth",
    "23e_0.9904acc.pth"
]
M_list = [os.path.join(_p, _m) for _p, _m in zip(P_list, M_list)]

C_list = [
    "C.26e.npy",
    "C.25e.npy",
    "C.23e.npy",
    "C.19e.npy",
    "C.23e.npy"
]
C_list = [os.path.join(_p, _c) for _p, _c in zip(P_list, C_list)]


for _m, _c, _type in zip(M_list, C_list, T_list):
    model = torch.load(_m).cuda()
    C = torch.from_numpy(np.load(_c)).float().cuda()
    # substitute weight vector with centre
    model.clf_layer.weight.data = C
    acc = test(model, _type)
    log_file.write("{}: {} ---\n".format(_type, acc))


log_file.write("end time: {}\n".format(time.asctime()))
log_file.flush()
log_file.close()
